import httpx
import json

async def generate(headers, json_data):
    timeout = httpx.Timeout(60.0)
    content_text = ''
    
    async with httpx.AsyncClient(timeout=timeout) as client:
        try:
            async with client.stream('POST', 'https://api.netqianye.com/v1/chat/completions', 
                                    headers=headers, json=json_data) as response:
                async for line in response.aiter_lines():
                    if not line or not line.startswith('data: '):
                        continue
            
                    if line == 'data: [DONE]':
                        # 计算token
                        try:
                            # 计算提示词token
                            prompt_text = ""
                            # 处理所有消息，不只是第一条
                            for message in json_data['messages']:
                                # 处理多模态内容
                                if isinstance(message.get('content'), list):
                                    for item in message['content']:
                                        if item.get('type') == 'text':
                                            prompt_text += item.get('text', '')
                                else:
                                    # 处理纯文本内容
                                    prompt_text += str(message.get('content', ''))
                            
                            # token估算 - 中文约1字1token，英文约4字符1token
                            def estimate_tokens(text):
                                chinese_char_count = sum(1 for c in text if '\u4e00' <= c <= '\u9fff')
                                other_char_count = len(text) - chinese_char_count
                                return (
                                    chinese_char_count +  # 中文字符数量
                                    (other_char_count // 4) +  # 英文字符数量整除4
                                    (1 if other_char_count % 4 else 0)  # 英文字符余数(有则加1)
                                )
                            
                            prompt_tokens = max(1, estimate_tokens(prompt_text))
                            completion_tokens = max(1, estimate_tokens(content_text))
                            total_tokens = prompt_tokens + completion_tokens
                            
                            # 直接返回JSON
                            response_json = {"usage": {
                                "prompt_tokens": prompt_tokens,
                                "completion_tokens": completion_tokens,
                                "total_tokens": total_tokens
                            }, "finished": True}
                            yield json.dumps(response_json).encode() + b'\n'
                        except Exception as ex:
                            print(f"Token计算错误: {str(ex)}")
                            yield b'{"usage":{"prompt_tokens":1,"completion_tokens":1,"total_tokens":2},"finished":true}\n'
                        break
                    
                    try:
                        data = json.loads(line[6:])
                        content = data["choices"][0]["delta"].get("content", "")
                        if content:
                            content_text += content
                            # 直接返回JSON
                            response_json = {"content": content}
                            yield json.dumps(response_json).encode() + b'\n'
                    except Exception as e:
                        # 直接返回JSON
                        error_msg = str(e)
                        print(f"请求错误: {error_msg}")
                        response_json = {"error": error_msg, "finished": True}
                        yield json.dumps(response_json).encode() + b'\n'
        except Exception as e:
            print(f"请求错误: {e}")
            response_json = {"error": error_msg, "finished": True}
            yield json.dumps(response_json).encode() + b'\n'